# coding: utf8

from numpy import array, eye, zeros, trace, cov, linalg, argmax


def extract_rotation(source_cloud: array, target_cloud: array):
    """ Extract rotation by using Eigen vectors.
    :param source_cloud: the source point cloud.
    :param target_cloud: the target point cloud.
    :return: the rotation matrix.
    """
    # sigma is the covariance matrix.
    sigma = array(cov(source_cloud.T, target_cloud.T)[0:3, 3:6])

    # Construct Q sigma matrix.
    a_matrix = sigma - sigma.T
    q_sigma = zeros((4, 4))
    q_sigma[0, 0] = trace(sigma)
    q_sigma[0, 1:4] = q_sigma[1:4, 0] = [a_matrix[1, 2], a_matrix[2, 0], a_matrix[0, 1]]
    q_sigma[1:4, 1:4] = sigma + sigma.T - trace(sigma) * eye(3)

    w, v = linalg.eig(q_sigma)

    q0, q1, q2, q3 = v[:, argmax(w)]
    rotation_matrix = array([
        [q0 ** 2 + q1 ** 2 + q2 ** 2 + q3 ** 2, 2 * (q1 * q2 - q0 * q3), 2 * (q1 * q3 + q0 * q2)],
        [2 * (q1 * q2 + q0 * q3), q0 ** 2 + q2 ** 2 - q1 ** 2 - q3 ** 2, 2 * (q2 * q3 - q0 * q1)],
        [2 * (q1 * q3 - q0 * q2), 2 * (q2 * q3 - q0 * q1), q0 ** 2 + q3 ** 2 - q1 ** 2 - q2 ** 2]
    ])

    return rotation_matrix
